#!/usr/autodesk/maya/bin/mayapy
# -*- coding:utf-8 -*-

import os, sys, logging, optparse, re

import pymel.core as pc
sys.path.append(os.path.dirname(__file__))

import katana_component as kc
reload(kc)
import maya2scenegraphXML as maya2scenegraphXML
reload(maya2scenegraphXML)

logger = logging.getLogger('dsScenegraphXML_Exporter')

# Loading the plugins needed
pc.loadPlugin('AbcExport', quiet=True)


def scenegraph_export(*args):

    # Asset Library path
    assetPath = pc.Path('/mnt/jiuqian/project/film/bao1/production/asset/scenes')

    # Convert to Path Object of pymel
    shotPath = pc.Path(args[0][0])

    # Define the path to export
    exportPath = pc.Path(args[0][1])

    # Get filename removed extension name
    shotBaseName = shotPath.basename().rstrip(shotPath.ext)

    #
    cacheMode = args[0][2]

    if os.path.exists(shotPath) and os.path.isfile(shotPath):

        pc.openFile(shotPath, force=True)


    sf = int(pc.playbackOptions(min=True, q=True))

    ef = int(pc.playbackOptions(max=True, q=True))

    if len(args[0]) == 5:

        sf = int(args[0][3])

        ef = int(args[0][4])


    # print '----------------------------------------'
    # logger.info(get_sg_dict())

    export_xml(get_sg_dict(), sf, ef, shotBaseName, exportPath, assetPath, cacheMode)
    #
    if cacheMode == 'shot':

        # Export the proxy file only when the mode is "shot"
        export_proxy(sf, ef, shotBaseName, exportPath)


def export_xml(sgList, sf, ef, shotBaseName, exportPath, assetPath, cacheMode):

    chr_dict, prop_dict, bg_dict = sgList

    assetDir_map = {
        'c': 'chr',
        'p': 'prop',
        'b': 'bg'
    }

    for dic in [chr_dict, prop_dict]:

        for values in dic.values():

            maya2scenegraphXML.setAssembly(values['assembly'])

            maya2scenegraphXML.setComponent(values['component'])

    for dic in [bg_dict]:

        for values in dic.values():

            maya2scenegraphXML.setAssembly(values['assembly'])

            maya2scenegraphXML.setStaticComponent(values['component'])


    roots = chr_dict.keys() + prop_dict.keys() + bg_dict.keys()


    for k in roots:

        klfPath = '%s/%s/%s' %(assetPath,
                               assetDir_map[k[0]],
                               k.rstrip('_geometry_GRP'))

        maya2scenegraphXML.setArbAttr([k], 'lookfile', "%s/%s.klf" %
                                      (klfPath,
                                       k.rstrip('_geometry_GRP')),
                                      'string')

    for mod in chr_dict.values():

        for comp in mod['component']:

            maya2scenegraphXML.setProxy([comp], '%s_proxy.abc' %'_'.join(
                comp.split('_')[:-1]))

    for mod in prop_dict.values():

        for comp in mod['component']:

            maya2scenegraphXML.setProxy([comp], '%s_proxy.abc' % '_'.join(
                comp.split('_')[:-1]))

    for mod in bg_dict.values():

        for comp in mod['component']:

            maya2scenegraphXML.setProxy([comp], '%s_proxy.abc' % '_'.join(
                comp.split('_')[:-1]))

    if cacheMode == 'lookdev':

        # Change export path to "/mnt/cache/lookfile/ASSET_TYPE/ASSET_NAME/" when cacehMode is "lookdev",
        # and don't attach "lookfile" attribute

        if shotBaseName[0] in assetDir_map.keys():
            lookfile_dir = assetDir_map[shotBaseName[0]]
            exportPath = '%s/lookfile/%s/%s' %(exportPath, lookfile_dir, '_'.join(shotBaseName.split('_')[:3]))
        else:
            logger.error("%s is not a type of [\"charactor\", \"prop\", \"background\"]." %shotBaseName)
            return


        #TODO: 2018.03.07, modifi the structure of Cache
        # PREFIX_data
        #       -- scn (.xml, abc)
        #       -- cfx (ass, abc)
        #       -- vfx (abc)
        make_subDir(exportPath, shotBaseName)

        maya2scenegraphXML.maya2ScenegraphXML(roots, '%s/%s_data/%s.xml' % (exportPath, shotBaseName, shotBaseName),
                                              startFrame=0,
                                              endFrame=0)

    else:

        make_subDir(exportPath, shotBaseName)
        # shotBaseName = '_'.join(shotBaseName.split('_')[:2])
        maya2scenegraphXML.maya2ScenegraphXML(roots, '%s/%s_data/%s.xml' % (exportPath, shotBaseName, shotBaseName),
                    startFrame=sf,
                    endFrame=ef,
                    arbAttrs=['lookfile'])


def make_subDir(parentDir, cacheName):

    fullPath = '%s/%s_data' % (parentDir, cacheName)

    if not os.path.isdir(fullPath):
        aniDir = fullPath + '/scn'
        cfxDir = fullPath + '/cfx'
        vfxDir = fullPath + '/vfx'

        for d in [aniDir, cfxDir, vfxDir]:
            os.makedirs(d)



def export_proxy(sf, ef, shotBaseName, exportPath):

    H_RIG_SUFFIX = 'h.mb'

    L_RIG_SUFFIX = 'l.mb'

    P_RIG_SUFFIX = 'p.mb'

    P_GRP_SUFFIX = '_px'

    mayaRefNodes = pc.ls(type='reference')

    for fileRefNode in mayaRefNodes:

        # Not "sharedReferenceNode"
        if not fileRefNode.name().endswith('sharedReferenceNode') and fileRefNode.name() != '_UNKNOWN_REF_NODE_':

            # Convert to FileReference Node of pymel

            # is loaded
            if fileRefNode.isLoaded():

                sliceFilePath = ''

                referenceFile = fileRefNode.referenceFile()

                if referenceFile.path.endswith(H_RIG_SUFFIX):

                    # Get the path to Slice Rig version
                    sliceFilePath = referenceFile.path.replace(H_RIG_SUFFIX, P_RIG_SUFFIX)

                elif referenceFile.path.endswith(L_RIG_SUFFIX):

                    # Get the path to Slice Rig version
                    sliceFilePath = referenceFile.path.replace(L_RIG_SUFFIX, P_RIG_SUFFIX)

                if os.path.exists(sliceFilePath):

                    # Replace file reference with slice version
                    referenceFile.replaceWith(sliceFilePath)

    proxyFiles = []
    # for p in P_GROUP:
    #     proxyGroups.append([pc.ls('*:*_%s' %p), p])

    # [[ [grp1, grp2], 'slice_grp'], [grp3, grp4], 'slice_grp']]
    for proxyGroup in pc.ls('*:*%s' %P_GRP_SUFFIX):

        proxyFileName = proxyGroup.replace('%s' %P_GRP_SUFFIX, '_') + "proxy.abc"

        proxyFilePath = '%s/%s_data_abc/%s' % (exportPath, shotBaseName, proxyFileName)

        pc.AbcExport(j="-fr %d %d -uv -sn -root %s -file %s" % (sf, ef, proxyGroup, proxyFilePath))

        proxyFiles.append(proxyFilePath)

    for f in proxyFiles:

        print "Exporting Proxy: \033[4;90m%s\033[0m" %(f.split(os.path.sep)[-1])



def get_sg_dict():

    grp_lst = pc.ls(type='transform')

    asset_lst = [ast for ast in grp_lst if ast.hasAttr('asset_name')]

    #TODO, 2018.03.05, if Asset is not exist, throw error.
    assert len(asset_lst) > 0, 'There are not any asset in the scene.'

    asset_dict = {'c': {}, 'p': {}, 'b': {}}

    for ast in asset_lst:

        is_static = ast.hasAttr('sgxml_static')

        # print '-=-=-=-=-=-'
        # print '%s is %s' %(ast, str(is_static))

        asset_type = ast.asset_type.get()

        asset_name = ast.asset_name.get()

        geo_grp = ast.getChildren()[0]

        if not '_geometry_' in geo_grp.name():
            continue

        assembly_grp = geo_grp.getChildren()

        if len(assembly_grp):

            # Set a asset as a BG, if it has "sgXML_static" attribute.
            if is_static:

                asset_type = 'b'


            asset_dict[asset_type][geo_grp.name()] = {}

            asset_dict[asset_type][geo_grp.name()]['assembly'] = [ast_grp.name() for ast_grp in assembly_grp]

            for assembly in assembly_grp:

                cps = assembly.getChildren()

                if len(cps):

                    if asset_dict[asset_type][geo_grp.name()].has_key('component'):

                            asset_dict[asset_type][geo_grp.name()]['component'].extend([comp_grp.name() for comp_grp in cps])

                    else:

                        asset_dict[asset_type][geo_grp.name()]['component'] = [comp_grp.name() for comp_grp in cps]



    return asset_dict['c'], asset_dict['p'], asset_dict['b']




if __name__ == '__main__':

    if len(sys.argv) > 1:

        if os.path.exists(sys.argv[1]) and os.path.isfile(sys.argv[1]):

            import maya.standalone

            maya.standalone.initialize()

            # Set environment value for maya
            os.environ['jiuqian_project'] = '/mnt/jiuqian/project/film'
            os.environ['test_project'] = '/mnt/project/bao/Proj_Maya/new_project/scenes/lookfile_test'

            scenegraph_export(sys.argv[1:])

            maya.standalone.uninitialize()
    else:
        print '\033[4;90m-------------------------------------\033[0m'

        print '\033[90mScenegraphXML_Export 1.0.0 linux "DeltaStudio VFX"\033[0m\n'

        print '\033[90mUsage: dsScenegraphXML_export.py YOUR_MAYA_FILE\033[0m'

        print '\033[4;90m-------------------------------------\033[0m'